from django.shortcuts import render
from dvadmin.knowledge.models import Knowledge,KnowledgeRecord,KnowledgeRecordSegment,KnowledgeHit
from dvadmin.utils.serializers import CustomModelSerializer
from dvadmin.utils.viewset import CustomModelViewSet
from dvadmin.utils.json_response import SuccessResponse, ErrorResponse, DetailResponse
from django.core.files.storage import default_storage
from django.db import transaction
from rest_framework.response import Response
import os,random
from dvadmin.db.milvus import MilvusUtil
from dvadmin.utils.fileUtil import FileFactory
from dvadmin.embedding.embedding import Bge
from pymilvus import FieldSchema,DataType,CollectionSchema


from application import settings
from rest_framework.views import APIView

from rest_framework.response import Response
from rest_framework_extensions.cache.mixins import CacheResponseMixin 

from langchain_community.vectorstores import Chroma
from langchain_community.llms.tongyi import Tongyi
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.agents import initialize_agent, AgentType
import json,_thread
from django_redis import get_redis_connection


from .tasks import deal_file

# Create your views here.
class KnowledgeSerializer(CustomModelSerializer):
    """
    接口白名单-序列化器
    """

    class Meta:
        model = Knowledge
        fields = "__all__"
        read_only_fields = ["id"]


class KnowledgeViewSet(CustomModelViewSet):
    """
    接口白名单
    list:查询
    create:新增
    update:修改
    retrieve:单例
    destroy:删除
    """
    queryset = Knowledge.objects.all()
    serializer_class = KnowledgeSerializer
    # permission_classes = []
    
    def __init__(self,*args, **kwargs):
        self.milvusUtil= MilvusUtil()
        # pool = redis.ConnectionPool(host='127.0.0.1', port=6379, db=0)
        # conn = redis.StrictRedis(connection_pool=pool)
    
    def create(self, request, *args, **kwargs):
        serializer = self.get_serializer(data=request.data, request=request)
        serializer.is_valid(raise_exception=True)
        self.perform_create(serializer)
        
        # 向量数据库创建表
        self.createVector(serializer.data.get('id'))
        
        return DetailResponse(data=serializer.data, msg="新增成功")
    def destroy(self, request, *args, **kwargs):
        
        instance = self.get_object()
        # 删除向量数据库表
        tablename = "knowledge" + str(instance.id)
        self.milvusUtil.dropTable(tablename)
        instance.delete()
        
        return DetailResponse(data=[], msg="删除成功")
    
    def createVector(self, kid, *args, **kwargs):
        # fields = [
        #     FieldSchema(name="id", dtype=DataType.INT64, auto_id=True,is_primary=True),
        #     FieldSchema(name="kid", dtype=DataType.INT64, description="知识库ID"),
        #     FieldSchema(name="file_id", dtype=DataType.INT64, description="文件ID"),
        #     FieldSchema(name="file_valid", dtype=DataType.INT8, default_value=1, description="文件有效"),
        #     FieldSchema(name="chunk_valid", dtype=DataType.INT8, default_value=1, description="块有效"),
        #     FieldSchema(name="context", dtype=DataType.VARCHAR, max_length=500, description="分块内容"),
        #     FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR,dim=1536, description="vector")
        # ]
        # schema = CollectionSchema(fields, "这是一个文档知识库表")
        tablename = "knowledge" + str(kid)
        # self.milvusUtil.createTable(tablename,schema)
        schema = self.milvusUtil.create_schema()

        # 3.2. Add fields to schema
        schema.add_field(field_name="id", datatype=DataType.INT64, auto_id=False,is_primary=True)
        schema.add_field(field_name="kid", datatype=DataType.INT64)
        schema.add_field(field_name="file_id", datatype=DataType.INT64)
        schema.add_field(field_name="file_valid", datatype=DataType.INT8, default_value=1)
        schema.add_field(field_name="chunk_valid", datatype=DataType.INT8, default_value=1)
        schema.add_field(field_name="context", datatype=DataType.VARCHAR, max_length=5000)
        schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=1024)

        # 3.3 准备索引parameters
        index_params = self.milvusUtil.prepare_index_params()

        # 3.4 添加索引
        index_params.add_index(
            field_name="embedding",
            index_name="index_embedding",
            index_type="IVF_FLAT",
            metric_type="IP",
            params={ "nlist": 128 }
        )
        self.milvusUtil.createTableAndIndex(tablename,schema,index_params)

class KnowledgeRecordSerializer(CustomModelSerializer):
    """
    接口白名单-序列化器
    """
    class Meta:
        model = KnowledgeRecord
        fields = "__all__"
        read_only_fields = ["id"]


class KnowledgeRecordViewSet(CustomModelViewSet):
    """
    接口白名单
    list:查询
    create:新增
    update:修改
    retrieve:单例
    destroy:删除
    """
    queryset = KnowledgeRecord.objects.all()
    serializer_class = KnowledgeRecordSerializer
    # permission_classes = []
    
    # flag = False
    # knowledgeRecord = None
    
    # def __new__(cls):
    #     # 1. 判断类属性是否是空对象
    #     if cls.knowledgeRecord is None:
    #         cls.knowledgeRecord = super().__new__(cls)
    #     return cls.knowledgeRecord
    
    def __init__(self, suffix, basename, detail) -> None:
        
        # if KnowledgeRecordViewSet.flag:
        #     return
        
        # KnowledgeRecordViewSet.flag = True
            
        # super().__init__()
        # self.suffix = suffix
        # self.basename = basename
        # self.detail = detail
        
        # 初始化Redis连接
        self.conn = get_redis_connection("default")
        
        self.channel = 'knowledge'        # 定义频道
        
    
    def get_queryset(self):
        if self.action =='list':
            params = self.request.query_params
            parent = params.get('parent', None)
            print(params)
            if params:
                if parent:
                    queryset = self.queryset.filter(kid=parent)
                else:
                    queryset = self.queryset.filter(kid__isnull=True)
            else:
                queryset = self.queryset.filter(kid__isnull=True)
            return queryset
        else:
            return self.queryset
    
    def create(self, request, *args, **kwargs):
        file_path = ''
        filename = ""
        if request.FILES:
            file = request.FILES['file']
            filename = file.name
            # 生成文件路径  MEDIA_ROOT
            file_path = os.path.join(settings.BASE_DIR,settings.MEDIA_ROOT,filename) 
            with open(file_path, 'wb+') as destination:  
                for chunk in file.chunks():  
                    destination.write(chunk) 
        
        
        # 获取文件类型
        fileType = filename.split(".")[-1]
        fileUtil = FileFactory().get_sclass(fileType)
            
        # 文件加载
        docs = fileUtil.fileLoader(file_path)
        # 文件分割
        chunks = fileUtil.fileSplitChar(docs,'\n\n',600,0)
        # doc转成字符串
        chunkList = fileUtil.doc2List(chunks)
        
        kid = request.data.get('kid') 
        
        data = KnowledgeRecord.objects.create(
            kid = kid,
            name = filename,
            type = 1 if fileType == 'txt' else 2,
            num_char =  sum(map(lambda x: len(x), chunkList)),
            num_chunk = len(chunks),
            status = 2,
        )
        
        rid = data.id
        
         # 将分块数据存入sql数据库中，并获取存入后的每个id
        no_chunk = 0
        for item in chunkList:
            no_chunk = no_chunk + 1
            KnowledgeRecordSegment.objects.create(
                content=item,
                valid=1,
                num_char=len(item),
                kid=kid,
                rid=rid,
                no_chunk=no_chunk
            )
        
        ## 发布订阅消息
        # self.conn.publish(self.channel, rid)
        
        ## 异步任务处理
        print("异步处理"+str(rid))
        ##deal_file.delay(rid)
        deal_file(rid)
        
        return DetailResponse(msg="新增成功")
                 
            

# Create your views here.
class KnowledgeHitSerializer(CustomModelSerializer):
    """
    接口白名单-序列化器
    """

    class Meta:
        model = KnowledgeHit
        fields = "__all__"
        read_only_fields = ["id"]


class KnowledgeHitViewSet(CustomModelViewSet):
    """
    接口白名单
    list:查询
    create:新增
    update:修改
    retrieve:单例
    destroy:删除
    """
    queryset = KnowledgeHit.objects.all()
    serializer_class = KnowledgeHitSerializer
    # permission_classes = []
    
    def __init__(self, suffix, basename, detail) -> None:
        super().__init__()
        self.suffix = suffix
        self.basename = basename
        self.detail = detail        
        self.milvus = MilvusUtil()
    
    def create(self, request, *args, **kwargs):
        data= request.data
        tablename = "knowledge" + str(data.get('knowledge') ) 
        keyword = data.get('keyword')   
        topk = data.get('topk')         
        ret = self.milvus.search(tablename,keyword,topk,0.5,1)
        
        serializer = self.get_serializer(data=request.data, request=request)
        serializer.is_valid(raise_exception=True)
        self.perform_create(serializer)        
        
        ## hit明细表
        
        
        return DetailResponse(data=serializer.data, msg="新增成功")
    
from django.http import JsonResponse
from dashscope import Generation



class ChatRAG(APIView):
    # 认证：跳过JWT token认证
    authentication_classes = []
    permission_classes = []
    
    def __init__(self):
        self.milvus = MilvusUtil()
        self.embedding = Bge()
    
    def post(self, request):      
        
        data = json.loads(request.body)
        
        
        question = data.get('question')
        knowledge = data.get('knowledge')
        stream = data.get('stream')
        
        tablename = "knowledge" + str(knowledge)
        
        # 1. 获取问题向量
        question_embedding = self.embedding.embed_query(question)
        
        # 2. 获取问题相关的知识,相似检索
        rets = self.milvus.search1(tablename,[question_embedding],3)
        
        context = ''
        for ret in rets[0]:
            context += ret['entity']['context']+'\n'
        
        # 3. 构造prompt
        prompt =  "请根据下面的内容回答问题：\n\n{prompt_context}\n问题：{prompt_question}"
        promptTemplate = PromptTemplate.from_template(prompt)
        
        if (stream): 
        
        # promptTemplate.format(prompt_context=context,prompt_question=question)
        
            message = []
            
            message.append({"role": "user", "content": promptTemplate.format(prompt_context=context,prompt_question=question)})
            
            responses = Generation.call(
                model="qwen-turbo", 
                messages=message,
                stream=True,
                result_format="message",
                incremental_output=True
            )
            
            response = StreamingHttpResponse(self.stream_response(responses),content_type='text/event-stream')
            
            return response
        
            
        else:       
            llm = Tongyi()
            data = llm.invoke(promptTemplate.format(prompt_context=context,prompt_question=question))        
            return JsonResponse({'code':200,'data':data,'msg':'ok'})
        
        
    def stream_response(self,responses):
        for response in responses:
            data1 = response.output.choices[0]['message']['content']
            data = f"data: {data1}\n\n"
            if data1:
                yield data.encode('utf-8')  # 必须编码为字节串
            else:
                return "no message"    
        
    
    def get(self, request): 
        return request









from dashscope import Generation
from django.http import JsonResponse,StreamingHttpResponse

class BaseChatView(APIView):
    
    # 认证：跳过JWT token认证
    authentication_classes = []
    permission_classes = []
    
    def __init__(self):        
        self.milvus = MilvusUtil()
        self.embedding = Bge()
    def knowledgePrompt(self,question,knowledge):
        tablename = 'knowledge'+ str(knowledge)
        
        # 2、获取问题向量
        question_embedding = self.embedding.embed_query(question)
        
        # 3、获取问题相关的知识
        rets = self.milvus.search1(tablename,[question_embedding],5)
        
        # 4、构造prompt
        prompt =  "请根据下面的内容回答问题：\n\n{prompt_context}\n问题：{prompt_question}"
        
        context = ""
        for ret in rets[0]:
            context += ret.get("entity").get("context")+'\n'
        
        promptTemplate = PromptTemplate.from_template(prompt)
        ret = promptTemplate.format(prompt_context=context,prompt_question=question)
        print(ret)
        
        return ret
        
    def knowledgeQA(self,question,knowledge):
        tablename = 'knowledge'+ str(knowledge)
        
        # 2、获取问题向量
        question_embedding = self.embedding.embed_query(question)
        
        # 3、获取问题相关的知识
        rets = self.milvus.search1(tablename,[question_embedding],2)
        
        # 4、构造prompt
        prompt =  "请根据下面的内容回答问题：\n\n{prompt_context}\n问题：{prompt_question}"
        promptTemplate = PromptTemplate.from_template(prompt)
        
        context = ""
        for ret in rets[0]:
            context += ret.get("entity").get("context")+'\n'
        
        ret = promptTemplate.format(prompt_context=context,prompt_question=question)
        print(ret)
        
        llm = Tongyi()
        chain = LLMChain(
            llm=llm, 
            prompt=promptTemplate
        )

        answer = chain.invoke({'prompt_context': context ,'prompt_question': question})
        
        return answer.get("text")

class KGChatView(BaseChatView):
        
    def post(self,request):
        try:
            data = json.loads(request.body)
            # 现在你可以使用data里的数据了
            # 例如：your_processing_function(data)
        except json.JSONDecodeError:
            return JsonResponse({"code": 400, 'message': "Json 参数解析失败"})
                
        # 1、接收参数
        question = data.get('question')
        knowledge = data.get('knowledge')
        stream = data.get('stream')
        
        if (stream):
            prompt = self.knowledgePrompt(question,knowledge)
        
            ssemessages = []
            
            ssemessages.append({"role": "user", "content": prompt})
    
            responses = Generation.call(
                model="qwen-turbo",  
                messages=ssemessages,  
                result_format='message',
                stream=True,  # 设置输出方式为流式输出
                incremental_output=True  # 增量式流式输出
            ) 
            
            response = StreamingHttpResponse(
                self.generate_sse(responses),
                content_type="text/event-stream",
            )
            response["Cache-Control"] = "no-cache"
            return response     
        else:
            ret = self.knowledgeQA(question,knowledge)
            
            return JsonResponse({"code": 200, 'message': ret})
    
    def generate_sse(self,responses):
        for response in responses:
            data1 = response.output.choices[0]['message']['content']
            data = f"data: {data1}\n\n"
            if data1:
                yield data.encode('utf-8')  # 必须编码为字节串
            else:
                return "no message"
    
class ChatCallView(BaseChatView):
    """使用function_call 方法"""
    
    def __init__(self):       
        self.tools = []            
    
    # 封装模型响应函数
    def get_response(self,messages):
        response = Generation.call(
            model='qwen-turbo',
            messages=messages,
            tools=self.tools,
            seed=random.randint(1, 10000),  # 设置随机数种子seed，如果没有设置，则随机数种子默认为1234
            result_format='message'  # 将输出设置为message形式
        )
        return response
    
    def post(self,request):
        try:
            data = json.loads(request.body)
            # 现在你可以使用data里的数据了
            # 例如：your_processing_function(data)
        except json.JSONDecodeError:
            return Response({"code": 400, 'message': "Json 参数解析失败"})
                
        # 1、接收参数
        question = data.get('question')
        k_id = data.get('knowledge')
        
        # 获取知识库描述
        knowledge = Knowledge.objects.get(id=k_id)
        
        self.tools.append(                
            # 知识库查询
            {
                "type": "function",
                "function": {
                    "name": "get_knowledge",
                    "description": knowledge.description,
                    "parameters": {}  
                }
            }
        )
        
        messages = [
            {
                "content": question,
                "role": "user"
            }
        ]
        
        # 第一次调用大模型
        first_response = self.get_response(messages)
        assistant_output = first_response.output.choices[0].message
        
        # 如果模型判断无需调用工具，则将assistant的回复直接打印出来，无需进行模型的第二轮调用
        if 'tool_calls' not in assistant_output:  
            return Response({"code": 200, 'message':  assistant_output.content})
        else:
            # 获取工具调用
            ret = self.knowledgeQA(question,k_id)
            return Response({"code": 200, 'message':  ret})
        
    def get(self,request):
        return Response({"code": 200, 'message': "成功"})

class ChatView(APIView):
    
    # 认证：跳过JWT token认证
    authentication_classes = []
    permission_classes = []
    
    def post(self,request):
        try:
            data = json.loads(request.body)
            # 现在你可以使用data里的数据了
            # 例如：your_processing_function(data)
        except json.JSONDecodeError:
            return Response({"code": 400, 'message': "Json 参数解析失败"})
                
        # 1、接收参数
        question = data.get('question')
        knowledge = data.get('knowledge')
        
        agent = initialize_agent(
            MyAgentTool().tools(),
            Tongyi(),
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
        )
        # agent.run("2024年NBA冠军球队是哪只?")
        ret = agent.invoke(input={'input':question,'knowledge':knowledge})
        
        return Response({"code": 200, 'message':  ret.get('output')})
        
    def get(self,request):
        return Response({"code": 200, 'message': "成功"})
    


   
from langchain.agents import Tool, tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
from langchain.prompts import StringPromptTemplate
from langchain.chains import LLMChain
from typing import List, Union
from langchain.schema import AgentAction, AgentFinish, OutputParserException
import re
    #自定义工具类
class MyAgentTool:
    
    db = None
    def __init__(self) -> None:
        self.milvus = MilvusUtil()
        # self.embedding = Bge()
        
    def tools(self):
        return [
            Tool(
                name="knowledgeTool",
                description="当你需要回答关于教育行业大模型应用的问题时使用这个工具",
                func=self.run,
            )
        ]
    def run(self, question: str) -> str: 
        knowledge=9
        tablename = 'knowledge'+ str(knowledge)
        
        # 2、获取问题向量
        question_embedding = self.embedding.embed_query(question)
        
        # 3、获取问题相关的知识
        rets = self.milvus.search1(tablename,[question_embedding],5)
        
       # 4、构造prompt
        promptTemplate =  "请根据下面的内容回答问题：\n{prpmpt_context}问题：{prpmpt_question}"
        
        prompt = ""
        for ret in rets[0]:
            prompt += ret.get("entity").get("context")+'\n'
        
        # 5、使用chain调用大模型
        llm = Tongyi()
        chain = LLMChain(
            llm=llm, 
            prompt=PromptTemplate.from_template(promptTemplate)
        )

        answer = chain.invoke({'prpmpt_context': prompt ,'prpmpt_question': question})
        
        return answer.get("text")

# @tool("my_search_tool")
# def search(query: str) -> list[str]:
#     """通过搜索引擎查询"""
#     result = search_wrapper.search(query)
#     return [res["snippet"] for res in result]


    

